# Set up ----------------------
library(tidyverse)
library(lmtest)
library(multcomp)
library(sandwich)
library(car)
library(margins)
library(broom)

# load data 
yg <- read_rds("data/yougov.rds")
lucid <- read_rds("data/lucid.rds")

# Functions --------------------------
stata_robust <- function(model) {
  coeftest(model, vcov = vcovHC(model, type="HC2"))
}

# Compute quantities ------------------
# Cohen's d effect sizes 
lucid_effect_size <- 
  lucid %>% 
  group_by(is_inversion) %>% 
  summarise(
    avg = mean(legitimacy, na.rm = T),
    std_dev = sd(legitimacy, na.rm = T) ,
    n = n()
  ) 
lucid_pooled_sd <- 
  sqrt(
    ((lucid_effect_size$n[1]-1)*lucid_effect_size$std_dev[1]^2 + 
       (lucid_effect_size$n[2]-1)*lucid_effect_size$std_dev[2]^2) / 
      (lucid_effect_size$n[1] + lucid_effect_size$n[2] - 2)
  )
lucid_effect <- 
  (lucid_effect_size$avg[1] - lucid_effect_size$avg[2]) / lucid_pooled_sd

yg_effect_size <- 
  yg %>% 
  group_by(is_inversion) %>% 
  summarise(
    avg = mean(legitimacy, na.rm = T),
    std_dev = sd(legitimacy, na.rm = T) ,
    n = n()
  ) 
yg_pooled_sd <- 
  sqrt(
    ((yg_effect_size$n[1]-1)*yg_effect_size$std_dev[1]^2 + 
       (yg_effect_size$n[2]-1)*yg_effect_size$std_dev[2]^2) / 
      (yg_effect_size$n[1] + yg_effect_size$n[2] - 2)
  )
yg_effect <- 
  (yg_effect_size$avg[1] - yg_effect_size$avg[2]) / yg_pooled_sd


# Graphs ----------------
# Figure 1: means by margin
margin_avg_lucid <- 
  lucid %>% 
  group_by(treat_margin) %>% 
  summarise(
    avg = mean(legitimacy, na.rm = T),
    std_err = sd(legitimacy, na.rm = T) / sqrt(sum(!is.na(legitimacy))),
  ) %>%
  mutate(
    sample = "Lucid"
  )
margin_avg_yg <- 
  yg %>% 
  group_by(treat_margin) %>% 
  summarise(
    avg = mean(legitimacy, na.rm = T),
    std_err = sd(legitimacy, na.rm = T) / sqrt(sum(!is.na(legitimacy)))
  ) %>% 
  mutate(
    sample = "YouGov"
  )

margin_avg_lucid %>% 
  bind_rows(margin_avg_yg) %>%
  filter(!is.na(treat_margin)) %>%
  mutate(
    treat_margin = case_when(
      treat_margin == "w3" ~ 1,
      treat_margin == "w1" ~ 2,
      treat_margin == "l1" ~ 3,
      treat_margin == "l3" ~ 4,
      treat_margin == "l5" ~ 5
    )) %>% 
  ggplot(aes(x = treat_margin, y = avg)) + 
  geom_point(
    aes(shape = sample, color = sample),
    position = position_dodge(0.5)
  ) + 
  geom_errorbar(
    aes(ymin = avg-1.96*std_err, 
        ymax = avg+1.96*std_err, color = sample),
    width = 0, position = position_dodge(0.5)
  ) +
  geom_vline(xintercept = 2.5, linetype = "dashed", size = 0.3) + 
  ylim(1, 4) + 
  scale_x_continuous(
    limits = c(0.75, 5.25),
    breaks = seq(1, 5, by = 1),
    labels = c("Win 3%", "Win 1%", 
               "Lose 1%", "Lose 3%", "Lose 5%")
  ) +
  scale_color_manual(
    values = c("#0089a7", "#1b813e")
  ) +
  xlab(NULL) + ylab(NULL) + 
  theme_bw() +
  theme(legend.title = element_blank())
ggsave("output/figures/margins/means_by_margin.png",
       width = 5, height = 3)

# Figure 2: means by party, copartisan, and margin 
margin_party_co_avg_lucid <- 
  lucid %>% 
  group_by(treat_margin, is_gop, coparty) %>% 
  summarise(
    avg = mean(legitimacy, na.rm = T),
    std_err = sd(legitimacy, na.rm = T) / sqrt(sum(!is.na(legitimacy))),
  ) %>%
  mutate(
    sample = "Lucid"
  )
margin_party_co_avg_yg <- 
  yg %>% 
  group_by(treat_margin, is_gop, coparty) %>% 
  summarise(
    avg = mean(legitimacy, na.rm = T),
    std_err = sd(legitimacy, na.rm = T) / sqrt(sum(!is.na(legitimacy)))
  ) %>% 
  mutate(
    sample = "YouGov"
  )
margin_party_co_means <- 
  margin_party_co_avg_lucid %>% 
  bind_rows(margin_party_co_avg_yg) %>% 
  filter(!is.na(treat_margin) & !is.na(is_gop)) %>% 
  mutate(
    treat_margin = case_when(
      treat_margin == "w3" ~ 1,
      treat_margin == "w1" ~ 2,
      treat_margin == "l1" ~ 3,
      treat_margin == "l3" ~ 4,
      treat_margin == "l5" ~ 5
    ), 
    coparty = ifelse(coparty == 1, "Copartisan", "Opposition"),
    is_gop = ifelse(is_gop == 1, "Republican", "Democrat")
  )

margin_party_co_means %>% 
  ggplot(aes(x = treat_margin, y = avg)) + 
  geom_point(
    aes(shape = sample, color = sample),
    position = position_dodge(0.5)
  ) + 
  geom_errorbar(
    aes(ymin = avg-1.96*std_err, 
        ymax = avg+1.96*std_err, color = sample),
    width = 0, position = position_dodge(0.5)
  ) +
  geom_vline(xintercept = 2.5, linetype = "dashed", size = 0.3) +
  scale_x_continuous(
    limits = c(0.75, 5.25),
    breaks = seq(1, 5, by = 1),
    labels = c("Win 3%", "Win 1%", 
               "Lose 1%", "Lose 3%", "Lose 5%")
  ) +
  scale_color_manual(
    values = c("#0089a7", "#1b813e")
  ) +
  ylim(1, 4) + 
  xlab(NULL) + ylab(NULL) + 
  theme_bw() + 
  theme(legend.title = element_blank()) +
  facet_grid(vars(coparty), vars(is_gop))

ggsave("output/figures/margins/means_by_margin_party_coparty_wapo.png",
       width = 7, height = 5)


# Figure three: political knowledge 
lucid <- 
  lucid %>% 
  mutate(
    pol_know_bin = case_when(
      pol_know %in% c(0,1) ~ "low",
      pol_know %in% c(2,3) ~ "med",
      pol_know %in% c(4,5) ~ "hi",
      TRUE ~ NA_character_
    ) %>% 
      factor(levels = c("low", "med", "hi"))
  )

pol_know_gop <- 
  lm(legitimacy ~ treat_margin*is_gop*pol_know_bin + coparty +
       pol_interest + is_not_white + is_college + is_female + age_bin, 
     data = lucid)
stata_robust(pol_know_gop)

K <- 
  rbind(
    # margin -1
    c(0, 1, rep(0, 30)), # D low knowledge
    c(0, c(1, 0, 0), rep(0, 14), 1, rep(0, 13)), # D middle knowledge
    c(0, c(1, 0, 0), rep(0, 17), 1, rep(0, 10)),# D high knowledge
    c(0, c(1, 0, 0), rep(0, 11), c(1, 0, 0), rep(0, 14)), # R low knowledge
    c(0, c(1, 0, 0), rep(0, 11), c(1, 0, 0), c(1, 0, 0), rep(0, 5), c(1, 0, 0), rep(0, 3)), # R middle knowledge
    c(0, c(1, 0, 0), rep(0, 11), c(1, 0, 0), rep(0, 3), c(1, 0, 0), rep(0, 5), c(1, 0, 0)), # R high knowledge
    
    # margin -3
    c(0, 0, 1, rep(0, 29)), # D low knowledge
    c(0, c(0, 1, 0), rep(0, 15), 1, rep(0, 12)), # D middle knowledge
    c(0, c(0, 1, 0), rep(0, 18), 1, rep(0, 9)), # D high knowledge
    c(0, c(0, 1, 0), rep(0, 11), c(0, 1, 0), rep(0, 14)), # R low knowledge
    c(0, c(0, 1, 0), rep(0, 11), c(0, 1, 0), c(0, 1, 0), rep(0, 5), c(0, 1, 0), rep(0, 3)), # R middle knowledge
    c(0, c(0, 1, 0), rep(0, 11), c(0, 1, 0), rep(0, 3), c(0, 1, 0), rep(0, 5), c(0, 1, 0)), # R high knowledge
    
    # margin -5 
    c(0, 0, 0, 1, rep(0, 28)),# D low knowledge
    c(0, c(0, 0, 1), rep(0, 16), 1, rep(0, 11)), # D middle knowledge
    c(0, c(0, 0, 1), rep(0, 19), 1, rep(0, 8)), # D high knowledge
    c(0, c(0, 0, 1), rep(0, 11), c(0, 0, 1), rep(0, 14)), # R low knowledge
    c(0, c(0, 0, 1), rep(0, 11), c(0, 0, 1), c(0, 0, 1), rep(0, 5), c(0, 0, 1), rep(0, 3)), # R middle knowledge
    c(0, c(0, 0, 1), rep(0, 11), c(0, 0, 1), rep(0, 3), c(0, 0, 1), rep(0, 5), c(0, 0, 1)) # R high knowledge
  )
colnames(K) <- names(coef(pol_know_gop))
rownames(K) <- 
  c(
    paste0("l1_", rep(c("D_", "R_"), each = 3),
           rep(c("low", "med", "hi"), 2)),
    paste0("l3_", rep(c("D_", "R_"), each = 3),
           rep(c("low", "med", "hi"), 2)),
    paste0("l5_", rep(c("D_", "R_"), each = 3),
           rep(c("low", "med", "hi"), 2))
  )

pol_know_glht <- 
  glht(model = pol_know_gop, 
       vcov. = vcovHC(pol_know_gop, type="HC2"), 
       linfct = K)
pol_know_coef <- 
  tidy(pol_know_glht) %>% 
  mutate(
    margin = rep(c("Lose 1%", "Lose 3%", "Lose 5%"), each = 6) %>% 
      factor(levels = c("Lose 5%", "Lose 3%", "Lose 1%")),
    party = rep(c(rep("Democrats", 3), rep("Republicans", 3)), 3),
    `Political knowledge` = rep(c("Low", "Medium", "High"), 6) %>% 
      factor(levels = c("Low", "Medium", "High"))
  )

pol_know_coef %>% 
  ggplot(aes(x = margin, y = estimate)) + 
  geom_hline(yintercept = 0, linetype = "dashed", alpha = 0.8) + 
  geom_point(
    aes(color = `Political knowledge`, shape = `Political knowledge`),
    position = position_dodge(width = 0.2)) +
  geom_errorbar(
    aes(
      color = `Political knowledge`,
      ymin = estimate - 1.96 * std.error, 
      ymax = estimate + 1.96 * std.error
    ),
    width = 0.05,
    position = position_dodge(width = 0.2)
  ) + 
  xlab(NULL) + 
  ylab("Treatment effect on legitimacy") +
  coord_flip() + 
  facet_grid(cols = vars(party)) + 
  theme_bw()

ggsave("output/figures/margins/pol_know_treat.png",
       width = 6, height = 3)

# Models --------------------------
# Table 1
yg_mod <- 
  lm(legitimacy ~ treat_margin + coparty + 
       pol_interest + is_not_white + is_college + is_female + age_bin, 
     data = filter(yg, pid7 %in% c(1,2,3,5,6,7)))
lucid_mod <- 
  lm(legitimacy ~ treat_margin + coparty + 
       pol_interest + is_not_white + is_college + is_female + age_bin, 
     data = lucid)

stata_robust(yg_mod)
stata_robust(lucid_mod)

# Table 2: marginal effects 
yg_partisans <- filter(yg, pid7 %in% c(1,2,3,5,6,7))
yg_3interaction <- 
  lm(legitimacy ~ treat_margin*is_gop*coparty + 
       pol_interest + is_not_white + is_college + is_female + age_bin, 
     data = yg_partisans)
lucid_3interaction <- 
  lm(legitimacy ~ treat_margin*is_gop*coparty + 
       pol_interest + is_not_white + is_college + is_female + age_bin, 
     data = lucid)
yg_marg_rep <- margins(yg_3interaction, data = subset(yg_partisans, is_gop == 1))
yg_marg_dem <- margins(yg_3interaction, data = subset(yg_partisans, is_gop == 0))
lucid_marg_rep <- margins(lucid_3interaction, data = subset(lucid, is_gop == 1))
lucid_marg_dem <- margins(lucid_3interaction, data = subset(lucid, is_gop == 0))

